Skip to content

fix(executorch): copy the engine in bulk instead of byte by byte - #4473

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:trt-engine-serialization-perf
Open

fix(executorch): copy the engine in bulk instead of byte by byte#4473
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:trt-engine-serialization-perf

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

The problem

Serializing a TensorRT engine to bytes goes through bytes(tensor.untyped_storage()).
That reads like a bulk copy, but it iterates the storage one element at a time in Python,
at roughly two seconds per megabyte:

16.8 MB   bytes(untyped_storage())   36.27 s
16.8 MB   memoryview(... .numpy())    0.01 s

For a model whose graph produces many engines, this step dominates the entire export. A
program with 8.8 GB of engines spends about five hours here. One with 70 GB spends closer
to two days. It is easy to mistake for a slow engine build, because the export simply sits
there making no visible progress.

It is also wrong for a view

contiguous() is a no-op for a tensor that is already contiguous, including a row slice of
a larger tensor. In that case the storage holds neighbouring bytes, so the serialized
engine comes out longer than the tensor it came from:

base = torch.arange(8, dtype=torch.uint8).reshape(2, 4)
row = base[0]                                    # tensor is [0, 1, 2, 3]

bytes(row.contiguous().untyped_storage())        # [0, 1, 2, 3, 4, 5, 6, 7]
bytes(memoryview(row.contiguous().numpy()))      # [0, 1, 2, 3]

Today's engines happen to own their whole storage, so this does not bite in practice yet.
It is a trap for anyone who later produces the engine buffer as a slice.

The fix

Copy through a memoryview over the tensor's own buffer:

engine_info[ENGINE_IDX] = bytes(
    memoryview(serialized_engine.cpu().contiguous().view(torch.uint8).numpy())
)

The comment on the previous code explains what it was avoiding: .numpy().tobytes()
allocates a second full-size buffer, which roughly doubles peak memory for a multi-gigabyte
engine. memoryview keeps that property, since it is a view rather than a copy, while
doing the copy in one shot and respecting the tensor's bounds.

Testing

Verified byte-for-byte equality against the previous path for the shapes this code sees:

host contiguous uint8    old 8.4 MB   new 8.4 MB   identical
device uint8             old 8.4 MB   new 8.4 MB   identical
non-contiguous view      old 16.8 MB  new 8.4 MB   differs, and the new value is correct

The third row is the bug above: the old path returned the whole storage rather than the
tensor.

Also exported a multi-method model that produces ten engines totalling 8.8 GB. Before the
change the serialization step ran for over an hour and a half without finishing; after it,
the same step completes in minutes and produces a program that loads and runs.

@meta-cla meta-cla Bot added the cla signed label Aug 9, 2026
@github-actions github-actions Bot added the component: api [Python] Issues re: Python API label Aug 9, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 9, 2026 21:54
Serializing an engine to bytes went through bytes(tensor.untyped_storage()).
That looks like a bulk copy but iterates the storage one element at a time in
Python, at roughly two seconds per megabyte:

  16.8 MB   bytes(untyped_storage())   36.27 s
  16.8 MB   memoryview(... .numpy())    0.01 s

For a model whose graph produces many engines, that dominates the whole export.
A program with 8.8 GB of engines spends about five hours here, and one with
70 GB spends closer to two days.

Reading the storage is also wrong for a tensor that is a view. contiguous() is
a no-op for a row slice, so the storage carries neighbouring bytes and the
serialized engine comes out longer than the tensor:

  tensor [0, 1, 2, 3]
  bytes(untyped_storage()) -> [0, 1, 2, 3, 4, 5, 6, 7]
  memoryview(numpy())      -> [0, 1, 2, 3]

Copy through a memoryview over the tensor's own buffer. This keeps the property
the previous code was after, no extra full-size buffer for a large engine, while
copying in one shot and respecting the tensor's bounds.
@shoumikhin
shoumikhin force-pushed the trt-engine-serialization-perf branch from 21ac6db to 4063483 Compare August 9, 2026 23:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant